//
//  Task worker - design 4
//  Connects PULL socket to tcp://localhost:5557
//  Collects workloads from ventilator via that socket
//  Connects PUSH socket to tcp://localhost:5558
//  Sends results to sink via that socket
//
//  use goroutines, no poll
//
package main

import (
	"fmt"
	"os"
	"zmq"
	"strconv"
	"strings"
	"flag"
)

var (
	maxch = flag.Int("maxch", 100, "goroutines max")
)

func send_work(sender *zmq.Socket, work chan int) os.Error {
	for {
		w, ok := <-work
		if !ok {
			break
		}
		err := sender.Send([]byte(strconv.Itoa(w)), 0)
		if err != nil {
			return err
		}
	}
	return nil
}

func poll_work(items []zmq.PollItem, receiver *zmq.Socket, work chan int, done chan bool) os.Error {
	defer close(work)
	ch := make(chan int, *maxch)
	for {
		data, err := receiver.Recv(zmq.DONTWAIT)
		if err == zmq.ErrAgain {
			zmq.Sleep(1)
			select {
			case <-done:
				return nil
			default:
			}
			continue
		}
		if err != nil {
			return nil
		}
		ar := strings.Split(string(data), ":")
		if len(ar) != 2 {
			continue
		}
		fmt.Printf("%s.", ar)
		w, err := strconv.Atoi(ar[0])
		if err != nil {
			return nil
		}
		n, err := strconv.Atoi64(ar[1])
		if err != nil {
			return nil
		}
		ch <- 0
		go func(ch chan int, work chan int, w int, n int64) {
			zmq.Sleep(n * 1e6)
			work <- w
			<-ch
		}(ch, work, w, n)
	}
	return nil
}

func zmq_main() os.Error {
	context, err := zmq.Init(1)
	if err != nil {
		return err
	}
	defer context.Close()

	receiver, err := context.Socket(zmq.PULL)
	if err != nil {
		return err
	}
	defer receiver.Close()
	err = receiver.Connect("tcp://localhost:5557")
	if err != nil {
		return err
	}

	sender, err := context.Socket(zmq.PUSH)
	if err != nil {
		return err
	}
	defer sender.Close()
	err = sender.Connect("tcp://localhost:5558")
	if err != nil {
		return err
	}

	controller, err := context.Socket(zmq.SUB)
	if err != nil {
		return err
	}
	defer controller.Close()
	err = controller.Connect("tcp://localhost:5559")
	if err != nil {
		return err
	}
	err = controller.SetOptionByte(zmq.SUBSCRIBE, nil)
	if err != nil {
		return err
	}

	//  Process messages from receiver and controller
	items := make([]zmq.PollItem, 2)
	items[0].SetSocket(receiver)
	items[0].SetEvents(zmq.POLLIN)
	items[1].SetSocket(controller)
	items[1].SetEvents(zmq.POLLIN)

	exit := make(chan bool)
	done := make(chan bool)
	work := make(chan int)

	go func() {
		err := send_work(sender, work)
		if err != nil {
			errch <- err
		} else {
			exit <- true
		}
	}()

	go func(controller *zmq.Socket, work chan int, done chan bool) {
		_, err = controller.Recv(0)
		if err != nil {
			errch <- err
		} else {
			done <- true
		}
	}(controller, work, done)

	go func() {
		err := poll_work(items, receiver, work, done)
		if err != nil {
			errch <- err
		}
	}()
	<-exit
	return nil
}

var errch = make(chan os.Error)

func main() {
	fmt.Println("Task worker")
	exit := make(chan bool)
	go func() {
		for {
			err := <-errch
			if err != nil {
				fmt.Println(err)
				os.Exit(1)
			}
		}
	}()
	go func() {
		err := zmq_main()
		if err != nil {
			errch <- err
		} else {	
			exit <- true
		}
	}()
	<-exit
}
